from treelib import Tree
import os


def generate_target_tree(root_dir, target_folders, max_depth=4, output_file="目标目录树.txt"):
    # 初始化目录树（根节点为项目根目录，第 1 级）
    tree = Tree()
    root_node_id = "root"
    tree.create_node(os.path.basename(root_dir), root_node_id)

    # 遍历每个目标文件夹（app、router bp，属于第 2 级）
    for target in target_folders:
        target_path = os.path.join(root_dir, target)
        if not os.path.exists(target_path):
            print(f"警告：{target_path} 不存在，跳过")
            continue

        # 目标文件夹节点（第 2 级）
        target_node_id = os.path.abspath(target_path)
        tree.create_node(target, target_node_id, parent=root_node_id)

        # 计算目标文件夹的基础层级（根目录是 1，目标文件夹是 2）
        base_level = 2
        # 递归遍历目标文件夹，限制最大深度（base_level + 子目录深度 ≤ max_depth）
        for dirpath, dirnames, filenames in os.walk(target_path):
            # 排除不需要的目录（__pycache__、venv、.idea 等）
            dirnames[:] = [d for d in dirnames if d not in ["__pycache__", "venv", ".idea"]]

            # 计算当前目录的层级（相对于根目录）
            current_level = len(os.path.relpath(dirpath, root_dir).split(os.sep)) + 1
            # 超过 max_depth 则停止遍历当前分支
            if current_level > max_depth:
                dirnames[:] = []  # 清空子目录列表，终止 os.walk 继续向下遍历
                continue

            # 创建当前目录节点（确保唯一）
            current_dir_id = os.path.abspath(dirpath)
            parent_dir_id = os.path.abspath(os.path.dirname(dirpath))
            if not tree.contains(current_dir_id):
                tree.create_node(
                    os.path.basename(dirpath),
                    current_dir_id,
                    parent=parent_dir_id
                )

            # 只在 ≤ max_depth 的目录下创建文件节点（可选：删除可只显示目录）
            if current_level <= max_depth:
                for filename in filenames:
                    file_id = os.path.abspath(os.path.join(dirpath, filename))
                    if not tree.contains(file_id):
                        tree.create_node(filename, file_id, parent=current_dir_id)

    # 打印到控制台 + 保存到文件
    tree.show()


if __name__ == "__main__":
    ROOT_DIR = os.getcwd()  # 项目根目录（第 1 级）
    TARGET_FOLDERS = ["app"]  # 第 2 级
    MAX_DEPTH = 2  # 最大展开到第 4 级（可按需修改）

    generate_target_tree(ROOT_DIR, TARGET_FOLDERS, max_depth=MAX_DEPTH)
